package com.liqw;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.StringJoiner;

/**
 * 分布式限流器（支持动态调整并发数、断电恢复）
 * 版本：JDK 8 兼容，使用 StringJoiner 构建 Lua 脚本，避免直接使用 '+' 拼接。
 */
@Component
public class DistributedRateLimiter {

    // ==================== 常量定义 ====================
    private static final String PREFIX = "rate_limiter:";
    private static final String MAX_CONCURRENT_KEY_TEMPLATE = PREFIX + "max_concurrent:%s";
    private static final String CURRENT_COUNT_KEY_TEMPLATE = PREFIX + "current_count:%s";

    // ==================== 成员变量 ====================
    private final StringRedisTemplate stringRedisTemplate;

    private RedisScript<Long> acquireRedisScript;
    private RedisScript<Long> releaseRedisScript;

    // ==================== 构造函数 ====================
    /**
     * 构造函数，依赖注入 StringRedisTemplate
     * @param stringRedisTemplate Redis 操作模板
     */
    @Autowired
    public DistributedRateLimiter(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    /**
     * 在 Bean 初始化后执行，用于构建并预编译 Lua 脚本
     */
    @PostConstruct
    public void initScripts() {
        String acquireScriptContent = buildAcquireScript();
        String releaseScriptContent = buildReleaseScript();

        this.acquireRedisScript = new DefaultRedisScript<>(acquireScriptContent, Long.class);
        this.releaseRedisScript = new DefaultRedisScript<>(releaseScriptContent, Long.class);
    }

    /**
     * 使用 StringJoiner 构建获取许可的 Lua 脚本
     * @return 完整的 Lua 脚本字符串
     */
    private String buildAcquireScript() {
        StringJoiner joiner = new StringJoiner("\n");
        joiner.add("local current_key = KEYS[1]");
        joiner.add("local max_key = KEYS[2]");

        joiner.add("local max = redis.call('get', max_key)");
        joiner.add("if not max then");
        joiner.add("  return -1");
        joiner.add("end");

        joiner.add("max = tonumber(max)");
        joiner.add("local current = redis.call('get', current_key) or '0'");
        joiner.add("current = tonumber(current)");
        joiner.add("if current >= max then");
        joiner.add("    return 0");
        joiner.add("end");

        joiner.add("return redis.call('incr', current_key)");
        return joiner.toString();
    }

    /**
     * 使用 StringJoiner 构建释放许可的 Lua 脚本
     * @return 完整的 Lua 脚本字符串
     */
    private String buildReleaseScript() {
        StringJoiner joiner = new StringJoiner("\n");
        joiner.add("local current_key = KEYS[1]");

        joiner.add("local current = redis.call('get', current_key)");
        joiner.add("if not current then");
        joiner.add("    return 0");
        joiner.add("end");

        joiner.add("current = tonumber(current)");
        joiner.add("if current <= 0 then");
        joiner.add("    redis.call('del', current_key)");
        joiner.add("    return 0");
        joiner.add("end");

        joiner.add("if current == 1 then");
        joiner.add("    redis.call('del', current_key)");
        joiner.add("    return 0");
        joiner.add("end");

        joiner.add("return redis.call('decr', current_key)");
        return joiner.toString();
    }

    // ==================== 核心方法 ====================

    /**
     * 尝试获取并发许可（非阻塞）
     * @param limiterKey 限流器唯一标识（如：order_service、user_login）
     * @return true：获取成功，false：获取失败
     */
    public boolean tryAcquire(String limiterKey) {
        Assert.hasText(limiterKey, "限流器键(limiterKey)不能为空");

        String currentCountKey = String.format(CURRENT_COUNT_KEY_TEMPLATE, limiterKey);
        String maxConcurrentKey = String.format(MAX_CONCURRENT_KEY_TEMPLATE, limiterKey);

        List<String> keys = new ArrayList<>(2);
        keys.add(currentCountKey);
        keys.add(maxConcurrentKey);

        Long result = stringRedisTemplate.execute(acquireRedisScript, keys);
        return result != null && result > 0;
    }

    /**
     * 释放并发许可（必须在业务完成后调用，否则会导致并发数泄漏）
     * @param limiterKey 限流器唯一标识
     */
    public void release(String limiterKey) {
        Assert.hasText(limiterKey, "限流器键(limiterKey)不能为空");

        String currentCountKey = String.format(CURRENT_COUNT_KEY_TEMPLATE, limiterKey);
        stringRedisTemplate.execute(releaseRedisScript, Collections.singletonList(currentCountKey));
    }

    /**
     * 动态调整最大并发数
     * @param limiterKey        限流器唯一标识
     * @param newMaxConcurrent  新的最大并发数（必须 > 0）
     */
    public void updateMaxConcurrent(String limiterKey, int newMaxConcurrent) {
        Assert.hasText(limiterKey, "限流器键(limiterKey)不能为空");
        Assert.isTrue(newMaxConcurrent > 0, "最大并发数(newMaxConcurrent)必须大于0");

        String maxConcurrentKey = String.format(MAX_CONCURRENT_KEY_TEMPLATE, limiterKey);
        stringRedisTemplate.opsForValue().set(maxConcurrentKey, String.valueOf(newMaxConcurrent));
    }

    /**
     * 从业务数据恢复当前并发数（断电后或重启时调用）
     * @param limiterKey            限流器唯一标识
     * @param businessCurrentCount  业务系统统计的当前并发数（会被修正为 >= 0）
     */
    public void recoverCurrentCount(String limiterKey, int businessCurrentCount) {
        Assert.hasText(limiterKey, "限流器键(limiterKey)不能为空");

        String currentCountKey = String.format(CURRENT_COUNT_KEY_TEMPLATE, limiterKey);
        int correctedCount = Math.max(businessCurrentCount, 0);

        if (correctedCount == 0) {
            stringRedisTemplate.delete(currentCountKey);
        } else {
            stringRedisTemplate.opsForValue().set(currentCountKey, String.valueOf(correctedCount));
        }
    }

    /**
     * 关闭方法（空实现，保持接口一致性，便于未来扩展）
     */
    public void shutdown() {
        // No-op
    }
}